import sys
import base64
import os.path

from PySide6.QtWidgets import QWidget, QGridLayout, QLabel, QLineEdit, QPushButton, QMessageBox, QTextEdit, QFileDialog, QApplication, QComboBox


class PictureToBase64Widget(QWidget):

    def __init__(self):
        super().__init__()
        self._init()

    def _init(self):
        grid = QGridLayout()
        self.setLayout(grid)
        self.input_show_widget = QLabel('', self)
        self.input_widget = QPushButton('选择文件', self)
        self.output_widget = QTextEdit('', self)
        self.confirm = QPushButton('转换', self)

        x = 0
        grid.addWidget(QLabel('图片: ', self), x, 0)
        grid.addWidget(self.input_show_widget, x, 1, 1, 3)
        grid.addWidget(self.input_widget, x, 4)
        grid.addWidget(self.confirm, x, 5, 1, 2)
        x += 1
        grid.addWidget(QLabel('Base64编码: ', self), x, 0)
        grid.addWidget(self.output_widget, x, 1, 1, 6)

        for i in range(x + 1):
            grid.setRowStretch(i, 1)

        self.input_widget.clicked.connect(self._choice_input_file)
        self.confirm.clicked.connect(self._convert)

    def _choice_input_file(self):
        file_name, _ = QFileDialog.getOpenFileName(self, filter='All Files (*.jpg *.png)')
        if file_name:
            self.input_show_widget.setText(file_name)
            self.input_widget.setText('重新选择')

    def _convert(self):
        if not self.input_show_widget.text():
            QMessageBox.question(self, '错误', '请选择一个原始文件', QMessageBox.StandardButton.Ok)
            return
        input_path = self.input_show_widget.text()
        if not os.path.exists(input_path):
            QMessageBox.question(self, '错误', '文件不存在', QMessageBox.StandardButton.Ok)
            return
        suffix = input_path.split('.')[-1]
        with open(input_path, 'rb') as f:
            base64_content = str(base64.b64encode(f.read()), encoding='utf-8')
            base64_content = f'data:image/{suffix};base64,' + base64_content
            self.output_widget.setText(base64_content)


class Base64ToPictureWidget(QWidget):

    def __init__(self):
        super().__init__()
        self._init()

    def _init(self):
        grid = QGridLayout()
        self.suffixes = ['jpg', 'jpeg', 'png']
        self.setLayout(grid)
        self.input_widget = QTextEdit('', self)
        self.output_show_widget = QLineEdit('', self)
        self.output_suffix_widget = QComboBox(self)
        self.output_suffix_widget.addItems(self.suffixes)
        self.confirm = QPushButton('转换', self)

        x = 0
        grid.addWidget(QLabel('Base64编码: ', self), x, 0)
        grid.addWidget(self.input_widget, x, 1, 1, 4)
        grid.addWidget(self.confirm, x, 5)
        x += 1
        grid.addWidget(QLabel('图片保存地址: ', self), x, 0)
        grid.addWidget(self.output_show_widget, x, 1, 1, 3)
        grid.addWidget(QLabel('图片格式: ', self), x, 4)
        grid.addWidget(self.output_suffix_widget, x, 5)
        for i in range(x + 1):
            grid.setRowStretch(i, 1)

        self.confirm.clicked.connect(self._convert)

    def _convert(self):
        if not self.input_widget.toPlainText():
            QMessageBox.question(self, '错误', 'Base64 不能为空', QMessageBox.StandardButton.Ok)
            return
        if not self.output_show_widget.text():
            QMessageBox.question(self, '错误', '图片保存地址不能为空', QMessageBox.StandardButton.Ok)
            return
        base64_content = self.input_widget.toPlainText()
        if ',' in base64_content:
            base64_content = base64_content.split(',')[-1].strip()

        file_name = self.output_show_widget.text()
        if os.path.exists(file_name):
            if os.path.isdir(file_name):
                file_name = file_name + os.sep + base64_content.replace('/', '')[:16] + '.' + self.output_suffix_widget.currentText()
        else:
            parent_path_dir = os.path.dirname(file_name)
            if not os.path.exists(parent_path_dir):
                os.makedirs(parent_path_dir)
            has_suffix = False
            for suffix in self.suffixes:
                if file_name.endswith(suffix):
                    has_suffix = True
                    break
            if not has_suffix:
                file_name += '.' + self.output_suffix_widget.currentText()

        with open(file_name, 'wb') as f:
            f.write(base64.b64decode(base64_content))
        QMessageBox.question(self, '提示', '转换完成', QMessageBox.StandardButton.Ok)


if __name__ == '__main__':
    app = QApplication()
    win1 = PictureToBase64Widget()
    win = Base64ToPictureWidget()
    win1.show()
    win.show()
    sys.exit(app.exec())
